{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "96b7476a",
   "metadata": {},
   "source": [
    "# 10-714 Homework 4 Extension"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a89d459",
   "metadata": {},
   "source": [
    "This homework is an extension of homework 4, where you will be implementing the Transformer architecture. For this assignment, all the things you need to implement is in the file `python/needle/nn/nn_transformer.py`. Other things in the needle library remains the same. This homework extension is built on homework 4, so make sure to copy the solutions from homework 4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a5c18f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code to set up the assignment\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "%cd /content/drive/MyDrive/\n",
    "!mkdir -p 10714\n",
    "%cd /content/drive/MyDrive/10714\n",
    "!git clone https://github.com/dlsyscourse/hw4_extra.git\n",
    "%cd /content/drive/MyDrive/10714/hw4_extra\n",
    "\n",
    "!pip3 install --upgrade --no-deps git+https://github.com/dlsys10714/mugrade.git\n",
    "!pip3 install pybind11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c9fb467",
   "metadata": {},
   "outputs": [],
   "source": [
    "!make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45349235",
   "metadata": {},
   "outputs": [],
   "source": [
    "%set_env PYTHONPATH ./python\n",
    "%set_env NEEDLE_BACKEND nd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f54d7073",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('./python')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5945207",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the PTB dataset\n",
    "\n",
    "import urllib.request\n",
    "import os\n",
    "\n",
    "!mkdir -p './data/ptb'\n",
    "# Download Penn Treebank dataset\n",
    "ptb_data = \"https://raw.githubusercontent.com/wojzaremba/lstm/master/data/ptb.\"\n",
    "for f in ['train.txt', 'test.txt', 'valid.txt']:\n",
    "    if not os.path.exists(os.path.join('./data/ptb', f)):\n",
    "        urllib.request.urlretrieve(ptb_data + f, os.path.join('./data/ptb', f))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cea5c0a",
   "metadata": {},
   "source": [
    "## Transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a2f639",
   "metadata": {},
   "source": [
    "In the previous homework you have implemented two sequence models, the Recurrent Neural Network, and Long Short-Term Memory. These models were once the state-of-the-art and default architecture choices on sequence modelling tasks, including language generation, until recently when the famous paper \"[Attention Is All You Need](https://arxiv.org/abs/1706.03762)\" (Vaswani et al. 2017) came out in 2017. Since then, Transformers, a model architecture introduced in the aforementioned paper, have become the standard and most performant class of model on language tasks. \n",
    "\n",
    "You will be implementing a Transformer in `python/needle/nn/nn_transformer.py`.\n",
    "\n",
    "Transformers are composed of three mains components that you will implement. \n",
    "1. A masked multi-head attention mechanism that adaptively focuses on different timesteps of a sequence. \n",
    "2. A residual block consisting of the attention layer followed by a two-layer neural network applied independently at each timestep. \n",
    "3. A Transformer model consisting of several stacked residual blocks (in this homework you will implement a decoder-only transformer).\n",
    "\n",
    "![model](https://miro.medium.com/v2/1*ZCFSvkKtppgew3cc7BIaug.png)\n",
    "\n",
    "The above is a photo of the Transformer architecture from Vaswani et al. 2017. The version of the transformer you will implement is nearly identical, but has layer normalization applied at the start of each residual block (referred to as a [prenorm variant](https://arxiv.org/abs/2002.04745) of the Transformer)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f094ff30",
   "metadata": {},
   "source": [
    "## Part 1: Implementing the Multi-Head Attention Activation Layer\n",
    "\n",
    "In this subproblem, you will be implementing the `forward` function of a \"base\" attention activation layer `MultiHeadAttention` in `python/needle/nn/nn_transformer.py`. This activation layer will take in three inputs: \n",
    "<p style=\"text-align: center;\">multi-head queries $Q \\in R^\\mathcal{B \\times H \\times T \\times D}$, keys $K \\in R^\\mathcal{B \\times H \\times T \\times D}$, and values $V \\in R^\\mathcal{B \\times H \\times T \\times D}$</p>\n",
    "\n",
    "where $B$ is the batch size, $H$ is the number of attention heads, $T$ is the sequence length, and $D$ is the hidden dimension. \n",
    "\n",
    "The attention output $X \\in R^{B \\times H \\times T \\times D}$ is computed as follows:\n",
    "\n",
    "<p style=\"text-align: center;\">$X = \\text{softmax}(\\frac{Q K^T}{\\sqrt{D}}) V$</p>\n",
    "\n",
    "Note that the matrix multiplications above are batched. This functionality is not natively supported in needle yet, so we have provided a convenient function `matmul` for batched matrix multiplications in `MultiHeadAttention`. Your goal in this section is to return $X$ given the input queries, keys, and values. \n",
    "\n",
    "For auto-regressive Transformer, this attention should support causal masking using the function `self.create_causal_mask` we have provided. This is to make sure that the prediction of next token only depends on it's previous tokens. Specifically, causal masking is applying a mask before the softmax so that the softmax probability is computed over a masked matrix of $\\frac{Q K^T}{\\sqrt{D}}$. \n",
    "\n",
    "In addition, your implementation should apply dropout to the attention softmax $\\text{softmax}(\\frac{Q K^T}{\\sqrt{D}})$. You can use the `self.dropout` function of the `MultiHeadAttention` module.\n",
    "\n",
    "Importantly, this layer is only an activation function, and has no trainable variables (these come later).\n",
    "\n",
    "Once you have finished your implementation, test your code with the following test cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7eeaa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -l -v -k \"attention_activation\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d19da8e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m mugrade submit \"YOUR_KEY_HERE\" -k \"attention_activation\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e65aea6",
   "metadata": {},
   "source": [
    "## Part 2 Implementing the Self-Attention Layer with trainable parameters\n",
    "\n",
    "In this subproblem, you will use the `MultiHeadAttention` class you just implemented, and wrap it in a subclass of `Module` called `AttentionLayer` in `python/needle/nn/nn_transformer.py`. \n",
    "\n",
    "This layer implements the self-attention with prenorm (when k, and v are None in the `self.forward` call) and cross-attention (when k and v are present in the `self.forward` call). We have provided skeleton code with the appropriate layer attributes defined. Your job is to write the forward pass of the `AttentionLayer`. Note that you are implementing multi-head attention, where the number of attention heads is given by the `self.num_head` attribute of the `AttentionLayer` class.\n",
    "\n",
    "Given inputs $Q \\in R^\\mathcal{B \\times T \\times D'}$, keys $K \\in R^\\mathcal{B \\times T \\times D'}$, and values $V \\in R^\\mathcal{B \\times T \\times D'}$ where $B$ is the batch size, $T$ is the sequence length, and $D'$ is the embedding dimension. This layer performs the following computation sequentially:\n",
    "\n",
    "(1) map queries, key, and values to heads.\n",
    "\n",
    "<p style=\"text-align: center;\">$Q' = \\text{LayerNorm}_q (Q) \\; W_q$</p>\n",
    "\n",
    "<p style=\"text-align: center;\">$K' = \\text{LayerNorm}_k (K) \\; W_k$</p>\n",
    "\n",
    "<p style=\"text-align: center;\">$V' = \\text{LayerNorm}_v (V) \\; W_v$</p>\n",
    "\n",
    "where $\\text{LayerNorm}_q , \\text{LayerNorm}_k, \\text{LayerNorm}_v $ are the prenorm `self.prenorm_q`, `self.prenorm_k` and `self.prenorm_v` respectively.\n",
    "\n",
    "(2) unravel heads from the channels axis.\n",
    "\n",
    "<p style=\"text-align: center;\">$Q' \\in R^{B \\times T \\times (HD)} \\to Q' \\in R^{B \\times H \\times T \\times D} $</p>\n",
    "\n",
    "<p style=\"text-align: center;\">$K' \\in R^{B \\times T \\times (HD)} \\to K' \\in R^{B \\times H \\times T \\times D} $</p>\n",
    "\n",
    "<p style=\"text-align: center;\">$V' \\in R^{B \\times T \\times (HD)} \\to V' \\in R^{B \\times H \\times T \\times D} $</p>\n",
    "\n",
    "where $H$ and $D$ are `self.num_head` and `self.head_dim` respectively.\n",
    "\n",
    "(3) compute the multi-head attention activation.\n",
    "\n",
    "<p style=\"text-align: center;\">$X = \\text{softmax}(\\frac{Q' (K')^T}{\\sqrt{D}}) V'$</p>\n",
    "\n",
    "<p style=\"text-align: center;\">$X \\in R^{B \\times H \\times T \\times D} \\to X \\in R^{B \\times T \\times H \\times D} $</p>\n",
    "\n",
    "<p style=\"text-align: center;\">$X \\in R^{B \\times T \\times H \\times D} \\to X \\in R^{B \\times T \\times (HD)}$</p>\n",
    "\n",
    "The last two steps do a transpose and then reshape to get the hidden states to be the correct shape.\n",
    "\n",
    "(4) project back to the input space of the layer with `self.out_projection`\n",
    "\n",
    "<p style=\"text-align: center;\">$X' = X \\; W_o$</p>\n",
    "\n",
    "Your goal in this part is to return $X$ in the `self.forward` call of `AttentionLayer`. For debugging, you may capture the `probs` variable returned by the inner `MultiHeadAttention` module and store it in an attribute such as `self.probs` of the attention layer.\n",
    "\n",
    "Once finished, you may test your layer with the following test cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b2fe04",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -l -v -k \"attention_layer\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d0bfad",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m mugrade submit \"YOUR_KEY_HERE\" -k \"attention_layer\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fa8fb30",
   "metadata": {},
   "source": [
    "## Part 3 Implementing a prenorm residual Transformer Layer\n",
    "\n",
    "You now have all the parts necessary to build a full Transformer by this point. In this subproblem, you will assemble the attention layer with a feedforward network into a stackable residual block. We have provided starter code in the `TransformerLayer` class. \n",
    "\n",
    "You will need to define the necessary class attributes in the `self.__init__` call of the module `TransformerLayer`, and fill in the forward pass in `self.forward`. Your transformer layer should support dropout applied to $X'$ from the previous step before adding a residual connection. Implement the following pseudocode of the layer, properly handling the intermediate tensor shapes:\n",
    "\n",
    "x - current sequence of hidden states\n",
    "\n",
    "<p style=\"text-align: center;\">$x = x + \\text{Dropout}(\\text{Attention}(x))$</p>\n",
    "<p style=\"text-align: center;\">$x = x + \\text{Dropout}(\\text{Linear}_{2}(\\text{Dropout}(\\text{ReLU}(\\text{Linear}_{1}(\\text{LayerNorm1d}(x))))))$</p>\n",
    "\n",
    "For the MLP, there are two Linear layers $\\text{Linear}_{1}$ and $\\text{Linear}_{2}$:\n",
    "- $\\text{Linear}_{1}$: input shape `q_features`, output shape `hidden_size`\n",
    "- $\\text{Linear}_{2}$: input shape `hidden_size`, output shape `q_features`\n",
    "\n",
    "Once finished, run the following test cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59e0fd87",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -l -v -k \"transformer_layer\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b74a6ecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m mugrade submit \"YOUR_KEY_HERE\" -k \"transformer_layer\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e78953",
   "metadata": {},
   "source": [
    "## Part 4 Implementing the Transformer model\n",
    "\n",
    "In this subsection, you will compose the residual transformer layers you implemented in the previous part to build the full Transformer model. Fill in the code in the `Transformer` class by defining a set of `num_layers` `TransformerLayer` modules with the appropriat parameters passed in from the parent `Transformer` class. Then, implement the `self.forward` call of the `Transformer`. \n",
    "\n",
    "As is, your current Transformer layers are permutation-invariant, and cannot tell which position each token is in the sequence. To break this symmetry, you will add a positional embedding to your Transformer.\n",
    "\n",
    "The original Transformer paper uses sinusoidal positional embeddings, and then adds to the input embeddings before the first `TransformerLayer`. These work well, but a more common strategy in modern Transformers is to learn the positional embeddings. \n",
    "\n",
    "To do this, you should use `needle.nn.Embedding`. In your Transformer implementation, create a learnable positional encoding using `needle.nn.Embedding` from homework 4, with `num_embeddings` set as `sequence_len`. Given an input sequence, you should create a tensor that has the timestep id of each token in the sequence (timesteps have increasing value, representing the position of a token in time), and use it like a word id. \n",
    "\n",
    "Last, add the created positional encoding to the input token embeddings before your transformer layers.\n",
    "\n",
    "Once complete, submit the following test cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec5fb0a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pytest -l -v -k \"transformer_model\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c897377",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m mugrade submit \"YOUR_KEY_HERE\" -k \"transformer_model\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "899683fc",
   "metadata": {},
   "source": [
    "Now, you can train a Transformer language model on the Penn Treebank dataset:\n",
    "\n",
    "Note: make sure to initialize a transformer model in the class `LanguageModel` of `apps/models.py`; also for Transformers, the final linear head `self.linear` should take in input dimension `embedding_size`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d118e5db",
   "metadata": {},
   "outputs": [],
   "source": [
    "import needle as ndl\n",
    "sys.path.append('./apps')\n",
    "from models import LanguageModel\n",
    "from simple_ml import train_ptb, evaluate_ptb\n",
    "\n",
    "device = ndl.cuda()\n",
    "corpus = ndl.data.Corpus(\"data/ptb\")\n",
    "train_data = ndl.data.batchify(corpus.train, batch_size=256, device=device, dtype=\"float32\")\n",
    "model = LanguageModel(20, len(corpus.dictionary), hidden_size=32, num_layers=1, seq_model='transformer', seq_len=20, device=device)\n",
    "train_ptb(model, train_data, seq_len=20, n_epochs=10, device=device, lr=0.003, optimizer=ndl.optim.Adam)\n",
    "evaluate_ptb(model, train_data, seq_len=20, device=device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
